import numpy as np
# from policy_value.policy_value_tf_2rows import PolicyValueNet
from policy_value.policy_value_tf_grids import PolicyValueNet

from mcts_alphaZero_player import MCTS_player
from collections import deque
import random
# from game.game_optimization_2rows import Game, Same_range_problems
from game.game_optimization_grids import Game, Same_range_problems

class TrainPipeline():

    def __init__(self, init_model=None):
        # params of the same_range_problem and game
        self.range_lower = -2 #变量区间下界
        self.range_upper = 2 #变量区间上界
        self.dim_num = 20 # 问题维度问题
        self.split_num = 5 #每一维度水平数（分成几个点），暂时都相同
        self.func = 'Rosenbrock'
        self.srq = Same_range_problems(
            lower=self.range_lower,
            upper=self.range_upper,
            dim_num=self.dim_num,
            split_num=self.split_num,
            func_name=self.func)
        self.game = Game(self.srq)

        # training params
        self.learn_rate = 2e-3
        self.lr_multiplier = 1.0
        self.buffer_size = 10000 #数据池数据
        self.data_buffer = deque(maxlen=self.buffer_size)
        self.batch_size = 256 # mini-batch size for training
        self.c_puct = 5
        self.n_playout = 400
        self.temp = 1.0 # 温度参数，控制exploration程度
        self.epochs = 2 #一个batch跑2次
        self.kl_targ = 0.02
        self.opt_batch_num = 1500
        self.check_freq = 50
        self.best_opt_val = float('-inf')
        # MCTS player构建
        self.policy_value_net = PolicyValueNet(self.dim_num,
                                               self.split_num,
                                               model_file=init_model)
        self.mcts_player = MCTS_player(self.policy_value_net.policy_value_fn,
                                       c_puct=self.c_puct,
                                       n_playout=self.n_playout,
                                       training=True)


    def collect_selfplay_data(self):
        """数据采集过程：collect self-play data for training"""
        train_data = self.game.start_self_play(self.mcts_player, temp=self.temp)
        # play_data = list(zip(*play_data))
        # train_data = [(state, probs, values) for state, probs, values in play_data]
        self.data_buffer.extend(train_data)

    def policy_evaluate(self, opt_times=10):
        mcts_player = MCTS_player(self.policy_value_net.policy_value_fn,
                                  c_puct=1,
                                  n_playout=self.n_playout,
                                  training=False)
        values = []
        states = []
        for i in range(opt_times):
            value_i, state_i = self.game.start_play(mcts_player)
            values.append(value_i)
            states.append(state_i)
        value_mean = np.mean(values) #平均价值
        states = np.array(states)
        return value_mean, states

    def policy_update(self):
        """update policy_value_net"""
        mini_batch = random.sample(self.data_buffer, self.batch_size)
        state_batch = [data[0] for data in mini_batch]
        mcts_probs_batch = [data[1] for data in mini_batch]
        value_batch = [data[2] for data in mini_batch]
        old_probs, old_v = self.policy_value_net.policy_value(state_batch)

        ##为防止policy大幅度变化，通过前后2次策略相似度KL值控制一组batchdata的训练次数与学习率
        # 控制训练次数
        for i in range(self.epochs):
            loss, entropy = self.policy_value_net.train_step(
                    state_batch,
                mcts_probs_batch,
                value_batch,
                self.learn_rate*self.lr_multiplier)
            new_probs, new_v = self.policy_value_net.policy_value(state_batch)
            kl = np.mean(np.sum(old_probs * (
                    np.log(old_probs + 1e-10) - np.log(new_probs + 1e-10)),
                    axis=1)
            )
            if kl > self.kl_targ * 4:  # early stopping if D_KL diverges badly
                break
        if kl > self.kl_targ * 2 and self.lr_multiplier > 0.1:
            self.lr_multiplier /= 1.5
        elif kl < self.kl_targ / 2 and self.lr_multiplier < 10:
            self.lr_multiplier *= 1.5

        # print(("kl:{:.5f},"
        #       "lr_multiplier:{:.3f}"
        #       "learn rate:{:.4f}"
        #       "loss:{:.4f},"
        #       "entropy:{}:.4f,").format(kl,
        #                                 self.lr_multiplier,
        #                                 self.learn_rate*self.lr_multiplier,
        #                                 loss,
        #                                 entropy))

        return loss, entropy

    def run(self):
        """running the training pipeline"""
        loss_list = []
        entrop_list = []
        for i in range(self.opt_batch_num):
            # 步骤1：数据采集
            self.collect_selfplay_data()

            # 步骤2：获取策略的loss，并且更新策略
            if len(self.data_buffer) > self.batch_size:
                loss, entropy = self.policy_update()
                print(("opt:{}  "
                       "loss:{:.4f}  ".format(i,
                                               loss)))
                loss_list.append(loss)
                entrop_list.append(entropy)
            # 步骤3：定期查看现有策略性能，判断是否保留参数
            if (i+1) % self.check_freq == 0:
                print("current self-play batch: {}".format(i + 1))
                opt_val, optimized_variables = self.policy_evaluate()
                self.policy_value_net.save_model('./current_policy.model')

                if opt_val > self.best_opt_val:
                    print("New best policy!!!!!!!!")
                    self.best_opt_val = opt_val
                    # update the best_policy
                    self.policy_value_net.save_model('./best_policy.model')

if __name__ == "__main__":
    training_pipeline = TrainPipeline(init_model="best_policy.model")
    training_pipeline.run()